import streamlit as st
import random
from PIL import Image
import tempfile
import pandas as pd
import time
import base64
import requests
import json


# 在 TrafficSimulator 类中添加API调用方法
# 修改后的TrafficSimulator类
class TrafficSimulator:
    def __init__(self):
        self.hospitals = pd.DataFrame({
            'name': ['上海仁济医院', '华山医院急诊部'],
            'distance': [2.5, 3.8],
            'beds': [1200, 800]
        })
        # API配置
        """
        self.main_api_url = ''
        self.main_api_key = ''
        self.main_model_name = 'Qwen/Qwen2.5-VL-72B-Instruct'
        """
    def _call_multimodal_api(self, messages):
        """通用API调用方法"""
        headers = {
            "Authorization": f"Bearer {self.main_api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.main_model_name,
            "messages": messages,
            "max_tokens": 2000
        }
        try:
            response = requests.post(
                self.main_api_url,
                headers=headers,
                json=payload,
                timeout=30
            )
            response.raise_for_status()
            return response.json()['choices'][0]['message']['content']
        except Exception as e:
            return f"API错误：{str(e)}"

    def _image_to_base64(self, image_path):
        """将图片转换为Base64编码"""
        if not image_path:
            return ""
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')

    def multimodal_analysis(self, text, image_path):
        """多模态分析（集成真实API）"""
        messages = [{
            "role": "user",
            "content": [
                {"type": "text", "text": text},
                {"type": "image_url", "image_url": {
                    "url": f"data:image/jpeg;base64,{self._image_to_base64(image_path)}"}
                }
            ]
        }] if image_path else [{
            "role": "user",
            "content": [
                {"type": "text", "text": text}
            ]
        }]
        return self._call_multimodal_api(messages)

    def generate_emergency_report(self, text, image_path):
        """生成应急预案报告"""
        prompt = f"""请根据以下情景生成应急预案和处理报告（Markdown格式）：
## 报告要求
1. 问题定级（1-10级）
2. 处理建议（分步骤说明）
3. 详细描述（包含可能的影响范围和持续时间）
4. 相关资源调度建议

## 输入信息
用户描述：{text}
"""
        messages = [{
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {
                    "url": f"data:image/jpeg;base64,{self._image_to_base64(image_path)}"}
                }
            ]
        }] if image_path else [{
            "role": "user",
            "content": [
                {"type": "text", "text": prompt}
            ]
        }]
        return self._call_multimodal_api(messages)

# ======== 界面系统 ========
st.set_page_config(
    page_title="Sky Eyes 模拟系统",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 初始化会话状态
if "conversations" not in st.session_state:
    st.session_state.conversations = {}
if "current_conv" not in st.session_state:
    st.session_state.current_conv = None
if "simulator" not in st.session_state:
    st.session_state.simulator = TrafficSimulator()

# 侧边栏对话管理
with st.sidebar:
    st.header("对话历史管理")

    # 新建对话按钮
    if st.button("➕ 新建对话"):
        new_id = str(int(time.time()))
        st.session_state.conversations[new_id] = {
            "id": new_id,
            "messages": [],
            "files": []
        }
        st.session_state.current_conv = new_id

    # 对话历史列表
    st.subheader("历史对话")
    for conv_id in list(st.session_state.conversations.keys())[::-1]:
        if st.button(
                f"对话 {conv_id[-4:]}",
                key=f"conv_{conv_id}",
                use_container_width=True
        ):
            st.session_state.current_conv = conv_id

# 主界面布局
tab1, tab2 = st.tabs(["实时对话分析", "应急响应"])

# 修改后的提交分析部分
with tab1:
    if st.session_state.current_conv:
        current_conv = st.session_state.conversations[st.session_state.current_conv]

        # 显示对话历史
        st.subheader("分析对话")
        for msg in current_conv["messages"]:
            with st.chat_message(msg["role"]):
                if msg["type"] == "user":
                    if msg["image"]:
                        st.image(Image.open(msg["image"]), width=300)
                    if msg["text"]:
                        st.markdown(f"**用户描述**: {msg['text']}")
                else:
                    st.markdown(msg["content"])
                    if "analysis" in msg:
                        st.success(f"分析结论: {msg['analysis']}")

        # 用户输入区域
        with st.form(key="multimodal_input"):
            col1, col2 = st.columns([2, 1])
            with col1:
                user_text = st.text_area("交通情况描述",
                                         "南北高架发生两车追尾事故",
                                         height=100)
            with col2:
                uploaded_image = st.file_uploader("上传现场图片",
                                                  type=["png", "jpg", "jpeg"])

            if st.form_submit_button("提交分析"):
                # 保存用户输入
                image_path = None
                if uploaded_image:
                    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
                        tmp_file.write(uploaded_image.getvalue())
                        image_path = tmp_file.name

                # 生成分析结果和应急预案
                with st.spinner("AI分析中..."):
                    analysis = st.session_state.simulator.multimodal_analysis(
                        user_text, image_path
                    )
                    emergency_report = st.session_state.simulator.generate_emergency_report(
                        user_text, image_path
                    )

                # 保存到对话历史
                current_conv["messages"].append({
                    "role": "user",
                    "type": "user",
                    "text": user_text,
                    "image": image_path
                })

                current_conv["messages"].append({
                    "role": "assistant",
                    "type": "system",
                    "content": "分析完成",
                    "analysis": analysis
                })

                current_conv["emergency_report"] = emergency_report
                st.rerun()

# 修改后的Tab2内容
with tab2:
    st.subheader("应急处理报告")
    if st.session_state.current_conv:
        current_conv = st.session_state.conversations[st.session_state.current_conv]

        if "emergency_report" in current_conv:
            st.markdown("### 📑 最新应急预案")
            st.markdown(current_conv["emergency_report"], unsafe_allow_html=True)

            # 添加报告操作按钮
            cols = st.columns([1, 1, 1])
            with cols[0]:
                if st.button("🖨️ 导出报告"):
                    st.download_button(
                        label="下载MD文件",
                        data=current_conv["emergency_report"],
                        file_name="emergency_report.md",
                        mime="text/markdown"
                    )
            with cols[2]:
                if st.button("🔄 重新生成"):
                    with st.spinner("重新生成报告中..."):
                        latest_input = next(
                            (msg for msg in reversed(current_conv["messages"])
                             if msg["role"] == "user"), None
                        )
                        if latest_input:
                            new_report = st.session_state.simulator.generate_emergency_report(
                                latest_input["text"], latest_input["image"]
                            )
                            current_conv["emergency_report"] = new_report
                            st.rerun()
        else:
            st.info("⚠️ 尚未生成应急预案，请先在实时分析页面提交分析请求")
    else:
        st.info("请先创建新对话")

